#!/usr/bin/env python3.9
# coding=utf-8

"""
Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
Description: python api test.
Author: MindX SDK
Create: 2022
History: NA
"""

import os
import logging
import sys
import time
from multiprocessing import Queue, Manager, Pool, Process
import traceback
import numpy as np
from mindx.sdk import base
from mindx.sdk.base import ImageProcessor
from mindx.sdk.base import VideoDecoder, VideoDecodeConfig, VdecCallBacker
from body import  *
from motor import *
from traffic import *
from portrait import *
from car_union import *
from car_plate import *
from config import *
import logging.config
import av
from sort import track_sort
logging.config.fileConfig("logger.conf")
logger = logging.getLogger("pipeline")


def video_decode(traffic_q: Queue, face_q: Queue, gpu_id: int, channel_id: int, video: dict):
    time_dict = dict()
    try:
        def vdec_callback(decode_image, channel_id, frame_id):
            # 1. Calculate preprocess resize and padding config
            decode_image.to_host()
            image_ori = np.array(decode_image.to_tensor())  # NHWC, RGB

            q_data = {
                "frame_id": frame_id,
                "image": image_ori,
                "height": decode_image.original_height,
                "width":  decode_image.original_width,
                "video_id": video["id"],
                "video_name": video["name"],
                "video_uri": video["uri"],
                "capture_time": common.current_time(),
            }

            traffic_q.put(q_data, block=True)
            face_q.put(q_data, block=True)

        # 1. initialize video decoder callback
        vdec_callbacker = VdecCallBacker()
        vdec_callbacker.registerVdecCallBack(vdec_callback)

        # 2. initialize video decoder config
        vdec_config = VideoDecodeConfig()
        vdec_config.skipInterval = SKIP_INTERVAL
        vdec_config.outputImageFormat = base.rgb
        vdec_config.width = VIDEO_WIDTH
        vdec_config.height = VIDEO_HEIGHT

        # 3. initialize video decode object
        video_decoder = VideoDecoder(vdec_config, vdec_callbacker, deviceId=gpu_id, channelId=channel_id)

        # 4. use ffmpeg av and VideoDecoder to decode video
        while True:
            with av.open(video["uri"], options={'rtsp_transport': 'tcp'}) as container:
                for frame_id, packet in enumerate(container.demux()):
                    if packet.size == 0:
                        continue
                    try:
                        video_decoder.decode(packet, frame_id)
                    except Exception as err:
                        logger.error("process0_video_decode decode error: ", err)
                    # if frame_id == VIDEO_BREAK_FRAME:
                    #    break
            # del video_decoder
            # 避免显存炸了
            time.sleep(10)
    except Exception as e:
        # video_decode(traffic_q, face_q, gpu_id, channel_id, video)
        print("process0_video_decode failed.", repr(e))
        traceback.print_exc()
        sys.exit(-1)


if __name__ == "__main__":
    base.mx_init()

    manager = Manager()

    # 视频解码 -> track
    q_decode_traffic = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_decode_face = manager.Queue(maxsize=MAX_QUEUE_SIZE)

    # 解码分发进程
    q_decode_traffic_list = list()
    q_decode_face_list = list()
    video_process_dict = dict()

    # 人像
    q_qua_face = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_feat_face = manager.Queue(maxsize=MAX_QUEUE_SIZE)

    # 人体
    q_rec_body = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_reid_body = manager.Queue(maxsize=MAX_QUEUE_SIZE)

    # 车辆
    q_rec_car_union_color = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_rec_car_union_prop = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_rec_car_union_brand = manager.Queue(maxsize=MAX_QUEUE_SIZE)
    q_det_rec_carplate = manager.Queue(maxsize=MAX_QUEUE_SIZE)

    # 非机动车
    q_rec_motor = manager.Queue(maxsize=MAX_QUEUE_SIZE)

    # 监控列表
    monitor_q_list = [
        ("待检测人机非", q_decode_traffic),
        ("待检测人脸", q_decode_face),
        ("人脸质量检测", q_qua_face),
        ("人脸解析", q_feat_face),
        ("人体解析", q_rec_body),
        ("人体重识别", q_reid_body),
        ("车辆颜色识别", q_rec_car_union_color),
        ("车辆属性识别", q_rec_car_union_prop),
        ("车辆品牌识别", q_rec_car_union_brand),
        ("车牌识别", q_det_rec_carplate),
        ("非机动车属性识别", q_rec_motor),
    ]

    # 统计模型加载数量
    model_loading = manager.list()

    video_list = []

    # 通过VAI接口获取rtsp/rtmp视频流列表,暂时写死
    if MODE == 'linewell':
        video_list = [
            {'id': '350500001', 'name': '点位1', 'uri': 'rtsp://admin:linewell123@10.0.1.233:554/ch1/main/'},
            {'id': '350500002', 'name': '点位2', 'uri': 'rtsp://wd:Aa123456@192.168.16.116:554/ch2/main/'},
        ]
    else:
        video_list = []
        with open(f"./data/{MODE}.txt", "r", encoding="UTF-8") as f:
            for line in f:
                point = line.split(',')
                video_list.append({'id': point[0], 'name': point[1], 'uri': f'rtsp://35.96.8.244:554/rtp/35000000002000000092_{point[0]}/'})

    for device_id in gpu_ids:
        # 1、检测：多进程，一张卡多进程
        for i in range(NUM_DETECTION_PROCESS):
            # 创建解码分发进程
            q_decode_traffic_dis = manager.Queue()
            q_decode_traffic_list.append(q_decode_traffic_dis)
            q_decode_face_dis = manager.Queue()
            q_decode_face_list.append(q_decode_face_dis)

            monitor_q_list.append((f"人像检测{device_id}_{i}", q_decode_traffic_dis))
            monitor_q_list.append((f"人机非检测{device_id}_{i}", q_decode_face_dis))
            # 人脸检测 -> 人脸track
            common.add_process(name=f"人像检测{device_id}_{i}", func=portrait_detect, args=(os.path.realpath(PORTRAIT_DET_MODEL_PATH), BATCH_SIZE, device_id, i, q_decode_face_dis, q_qua_face, model_loading))
            # 交通元素检测 -> 人体track、车辆track、非机动车track
            common.add_process(name=f"人机非检测{device_id}_{i}", func=traffic_detect, args=(os.path.realpath(TRAFFIC_DET_MODEL_PATH), BATCH_SIZE, device_id, i, q_decode_traffic_dis, q_rec_body, q_rec_car_union_color, q_rec_motor, model_loading))

        # 3：识别
        for i in range(NUM_PATTERN_RECOG_PROCESS):
            # 人脸质量
            common.add_process(name=f"人脸质量{device_id}_{i}", func=portrait_quality, args=(os.path.realpath(PORTRAIT_QUA_BB_MODEL_PATH), os.path.realpath(PORTRAIT_QUA_SC_MODEL_PATH), BATCH_SIZE, device_id, i, q_qua_face, q_feat_face, model_loading))

            # 人脸特征
            common.add_process(name=f"人脸特征{device_id}_{i}", func=portrait_feature, args=(os.path.realpath(PORTRAIT_FEAT_MODEL_PATH), BATCH_SIZE, device_id, i, q_feat_face, model_loading))

            # 人体属性识别
            common.add_process(name=f"人体属性识别{device_id}_{i}", func=body_recognize, args=(os.path.realpath(BODY_REC_MODEL_PATH), BATCH_SIZE, device_id, i, q_rec_body, q_reid_body, model_loading))

            # 车辆颜色识别
            common.add_process(name=f"车辆颜色识别{device_id}_{i}", func=car_union_color_recognize, args=(os.path.realpath(CARUNION_REC_COLOR_MODEL_PATH), CAR_BATCH, device_id, i, q_rec_car_union_color, q_rec_car_union_prop, model_loading))

            # 车辆属性识别
            common.add_process(name=f"车辆属性识别{device_id}_{i}", func=car_union_prop_recognize, args=(os.path.realpath(CARUNION_REC_PROP_MODEL_PATH), CAR_BATCH, device_id, i, q_rec_car_union_prop, q_rec_car_union_brand, model_loading))

            # 车辆品牌识别
            common.add_process(name=f"车辆品牌识别{device_id}_{i}", func=car_union_brand_recognize, args=(os.path.realpath(CARUNION_REC_BRAND_MODEL_PATH), CAR_BATCH, device_id, i, q_rec_car_union_brand, q_det_rec_carplate, model_loading))

            # 车牌检测和识别
            common.add_process(name=f"车牌检测和识别{device_id}_{i}", func=car_plate_detect_rec, args=(os.path.realpath(CARPLATE_DET_MODEL_PATH), os.path.realpath(CARPLATE_REC_MODEL_PATH), CAR_BATCH, device_id, i, q_det_rec_carplate, model_loading))

            # 非机动车属性识别
            common.add_process(name=f"非机动车属性识别{device_id}_{i}", func=motor_recognize, args=(os.path.realpath(BODY_REC_MODEL_PATH), BATCH_SIZE, device_id, i, q_rec_motor, model_loading))

        # 重识别
        for i in range(NUM_REID_PROCESS):
            common.add_process(name=f"人体重识别{device_id}_{i}", func=body_reid, args=(os.path.realpath(BODY_REID_MODEL_PATH), BATCH_SIZE, device_id, i, q_reid_body, model_loading))

    # 解码分发
    det_process_count = len(gpu_ids) * NUM_DETECTION_PROCESS

    while True:
        # 模型加载完成后再进行解码
        if len(model_loading) >= MODEL_NUM * len(gpu_ids):
            # 迭代rtsp视频流
            # 解码多进程，一路视频一个进程，一张卡25个进程
            for i, video in enumerate(video_list):
                gpu_id = gpu_ids[i % len(gpu_ids)]
                process_index = i % det_process_count
                logging.info(f"[decode]---->添加：{video['id']}, {video['name']},{video['uri']}")
                common.add_process(name=f"视频解码[{video['id']}]", func=video_decode, args=(q_decode_traffic_list[process_index], q_decode_face_list[process_index], gpu_id, i, video))
            break
        else:
            time.sleep(1)
            logging.info(f"Loading Model Num :{len(model_loading)} < {MODEL_NUM * len(gpu_ids)} ")

    # 暴力队列监控，当队列数量超过MAX_QUEUE_SIZE时清空队列，避免程序内存溢出
    common.add_process(name="监控进程", func=common.print_queues, args=(monitor_q_list, ))

    # 暴力进程监控，任何子进程，死了就重新拉起来
    common.monitor_processes()

    common.join_jobs()


